/* 
 * Copyright (c) 2012 Christoph Mueller <ruunhb@googlemail.com>
 * 
 * stal is free software: you can redistribute it and/or modify
 * it under the terms of the Lesser GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * stal is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * Lesser GNU General Public License for more details.
 *
 * You should have received a copy of the Lesser GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include "map.h"
#include <stdlib.h>
#include <assert.h>
#include <string.h>

enum RbNodeColor { RED, BLACK };

#define RB_PARENT(node) node->parent
#define RB_GRAND_PARENT(node) node->parent->parent

struct RbNode {
    const_pointer entry;
    const_pointer data;
    enum RbNodeColor color;
    struct RbNode* parent;
    struct RbNode* left;
    struct RbNode* right;
};


struct RbTree {
    struct RbNode* root;
    lq_ordering_funptr compare;
    size_t nodes;
};


struct RbTree*          lq_rbtree_new(lq_ordering_funptr compare);
void			lq_rbtree_clear(struct RbTree* tree, lq_free_funptr fun_entry, lq_free_funptr fun_data);

struct RbNode*          lq_rbtree_lookup(struct RbTree* tree, const_pointer entry);
struct RbNode*          lq_rbtree_insert(struct RbTree* tree, const_pointer entry, const_pointer data);
struct RbNode*          lq_rbtree_remove(struct RbTree* tree, const_pointer entry);

struct RbNode*          lq_rbnode_first(struct RbTree* tree);
struct RbNode*          lq_rbnode_last(struct RbTree* tree);

struct RbNode*          lq_rbnode_next(struct RbNode* node);
struct RbNode*          lq_rbnode_prev(struct RbNode* node);

struct RbNode*          lq_rbnode_new(enum RbNodeColor color, struct RbNode* parent, const_pointer entry, const_pointer data);
void			lq_rbnode_free(struct RbNode* node, lq_free_funptr fun_entry, lq_free_funptr fun_data);


struct LqMap*  
lq_map_new(lq_ordering_funptr compare)
{
	return lq_cast(struct LqMap*, lq_rbtree_new(compare));
}


void
lq_map_free(struct LqMap* map)
{
	lq_map_clear(map, 0, 0);

	free(map);
}


void 			
lq_map_clear(struct LqMap* map, lq_free_funptr fun_key, lq_free_funptr fun_data)
{
        assert(map != 0);
        lq_rbtree_clear(lq_cast(struct RbTree*, map), fun_key, fun_data);
}


size_t          
lq_map_size(struct LqMap* map)
{
	assert(map != 0);

	return lq_cast(struct RbTree*, map)->nodes;
}


pointer         
lq_map_lookup(struct LqMap* map, const_pointer key)
{
	assert(map != 0);

	struct RbNode* result = lq_rbtree_lookup(lq_cast(struct RbTree*, map), key);

	return (result != 0) ? lq_cast(pointer, result->data) : 0;
}


int             
lq_map_insert(struct LqMap* map, const_pointer key, const_pointer value)
{
	assert(map != 0);

	struct RbNode* result = lq_rbtree_insert(lq_cast(struct RbTree*, map), key, value);

	return result != 0;
}


struct LqPair
lq_map_remove(struct LqMap* map, const_pointer key)
{
	assert(map != 0);

	struct LqPair pair = { .key = 0, .data = 0 };

	struct RbNode* result = lq_rbtree_remove(lq_cast(struct RbTree*, map), key);

	if(result) {
		pair.key = result->entry;
		pair.data = result->data;

                lq_rbnode_free(result, 0, 0);
	}

	return pair;
}


struct LqPair*
lq_map_pairs(struct LqMap* map)
{
        assert(map != 0);

        size_t nodes = lq_map_size(map);

        if(nodes != 0) {
                struct RbNode* node = lq_rbnode_first(lq_cast(struct RbTree*, map));
                struct LqPair* pair = calloc(sizeof(struct LqPair), nodes);
                struct LqPair* begin = pair;
                
                while(node != 0) {
                        pair->key = node->entry;
                        pair->data = node->data;

                        pair++;
                        node = lq_rbnode_next(node);
                }

                return begin;
        }

        return 0;
}


struct RbTree* 
lq_rbtree_new(lq_ordering_funptr compare)
{
	assert(compare != 0);

	struct RbTree* tree = malloc(sizeof(struct RbTree));
	tree->root = 0;
	tree->compare = compare;
	tree->nodes = 0;

	return tree;
}


struct RbNode*
lq_rbnode_new(enum RbNodeColor color, struct RbNode* parent, const_pointer entry, const_pointer data)
{
	assert(entry != 0);

	struct RbNode* node = malloc(sizeof(struct RbNode));
	node->color = color;
	node->parent = parent;
	node->entry = entry;
	node->data = data;
	node->left = 0;
	node->right = 0;

	return node;
}


static void
lq_rbsubtree_free(struct RbNode* node, lq_free_funptr fun_entry, lq_free_funptr fun_data)
{
	if(node == 0)
		return;

	lq_rbsubtree_free(node->left, fun_entry, fun_data);
	lq_rbsubtree_free(node->right, fun_entry, fun_data);

	lq_rbnode_free(node, fun_entry, fun_data);
}


void
lq_rbtree_clear(struct RbTree* tree, lq_free_funptr fun_entry, lq_free_funptr fun_data)
{
	assert(tree != 0);

	lq_rbsubtree_free(tree->root, fun_entry, fun_data);
        
        tree->root = 0;
        tree->nodes = 0;
}


void			
lq_rbnode_free(struct RbNode* node, lq_free_funptr fun_entry, lq_free_funptr fun_data)
{
        assert(node != 0);

        if(fun_entry != 0)
                fun_entry(lq_cast(pointer, node->entry));

	if(fun_data != 0)
		fun_data(lq_cast(pointer, node->data));

        free(node);
}



struct RbNode*
lq_rbtree_lookup(struct RbTree* tree, const_pointer entry)
{
	assert(tree != 0 && entry != 0);

	struct RbNode* node = tree->root;

	while(node != 0) {
		int result = tree->compare(entry, node->entry);

		if(result < 0) 
			node = node->left;
		else if(result > 0)
			node = node->right;
		else
			return node;
	}

	return 0;
}



static struct RbNode*
lq_rbnode_rotate_left(struct RbNode* node)
{
	assert(node != 0);

	struct RbNode* parent = node->parent;
	struct RbNode* right = node->right;

	// switch right-node and its parent in a left-in-order rotation
	node->right = right->left;
	right->left = node;

	// correct parent hierarchy
	node->parent = right;
	right->parent = parent;
	
	if(node->right != 0) 
		node->right->parent = node;

	if(parent != 0 && parent->left == node)
		parent->left = right;
	else if(parent != 0 && parent->right == node)
		parent->right = right;

	// return new sub_root
	return right;
}


static struct RbNode*
lq_rbnode_rotate_right(struct RbNode* node)
{
	assert(node != 0);

	struct RbNode* parent = node->parent;
	struct RbNode* left = node->left;

	// switch left-node and its parent in a right-in-order rotation
	node->left = left->right;
	left->right = node;
	
	// correct parent hierarchy
	node->parent = left;
	left->parent = parent;

	if(node->left != 0)
		node->left->parent = node;

	if(parent != 0 && parent->left == node)
		parent->left = left;
	else if(parent != 0 && parent->right == node)
		parent->right = left;

	// return new sub_root
	return left;
}



static struct RbNode* 
lq_rbtrinode_restructering(struct RbNode* node)
{
	assert(node != 0);

	struct RbNode* root = RB_GRAND_PARENT(node);
	if(root->left == RB_PARENT(node)) {
		if(RB_PARENT(node)->left == node) 
			root = lq_rbnode_rotate_right(root);
		else {
			lq_rbnode_rotate_left(RB_PARENT(node));
			root = lq_rbnode_rotate_right(root);
		}

  	} else {
		if(RB_PARENT(node)->right == node)
			root = lq_rbnode_rotate_left(root);
		else {
			lq_rbnode_rotate_right(RB_PARENT(node));
			root = lq_rbnode_rotate_left(root);
		}
	}

	return root;
}


static void
lq_rbremedy_double_red(struct RbTree* tree, struct RbNode* node_z)
{
	assert(tree != 0 && node_z != 0);

	if(tree->root == RB_PARENT(node_z))
		return;
	else if(RB_PARENT(node_z)->color == BLACK)
		return;

	if(RB_GRAND_PARENT(node_z)->left == 0 || RB_GRAND_PARENT(node_z)->right == 0 || RB_GRAND_PARENT(node_z)->left->color == BLACK || RB_GRAND_PARENT(node_z)->right->color == BLACK) {
		// trinode restructering with single/multiple left- and right rotations
		struct RbNode* node_v = lq_rbtrinode_restructering(node_z);

		node_v->color = BLACK;
		node_v->left->color = RED;
		node_v->right->color = RED;

		while(node_v->parent != 0)
			node_v = node_v->parent;

		tree->root = node_v;
	} else {
		// recoloring of nodes
		RB_GRAND_PARENT(node_z)->left->color = BLACK;
		RB_GRAND_PARENT(node_z)->right->color = BLACK;

		if(RB_GRAND_PARENT(node_z)->parent == 0)
			return;

		RB_GRAND_PARENT(node_z)->color = RED;
		lq_rbremedy_double_red(tree, RB_GRAND_PARENT(node_z));
	}
}


static void
lq_rbremedy_double_black(struct RbTree* tree, struct RbNode* node_x, struct RbNode* node_r)
{
	assert(node_x != 0);

	struct RbNode* sibling = (node_x->left == node_r) ? node_x->right : node_x->left;
	struct RbNode* child = 0;

        if(sibling == 0)
                return;

	enum RbNodeColor parent_color = node_x->color;

	if(sibling->color == BLACK) {
		// get red child of siblings if available
		if(sibling->left != 0 && sibling->left->color == RED)
			child = sibling->left;
		else if(sibling->right != 0 && sibling->right->color == RED)
			child = sibling->right;

		if(child != 0) {
			// Case 1: Restructering of siblings red child
			child = lq_rbtrinode_restructering(child);

			child->color = parent_color;
                        if(node_r != 0)
        			node_r->color = BLACK;
			child->left->color = BLACK;
			child->right->color = BLACK;

			while(child->parent != 0)
				child = child->parent;
			tree->root = child;
		} else {
			// Case 2: Both siblings childs are black => Recoloring
                        if(node_r != 0)
			        node_r->color = BLACK;
			sibling->color = RED;
			if(node_x->color == RED)
				node_x->color = BLACK;
			else if(node_x->parent != 0) {
				lq_rbremedy_double_black(tree, RB_PARENT(node_x), node_x);
			}

		}
	} else {
		// Case 3: Sibling is red => Adjustment needed
		if(node_x->right == sibling)
			child = sibling->right;
		else
			child = sibling->left;

                if(child != 0) {
                        child = lq_rbtrinode_restructering(child);
                        sibling->color = BLACK;
                        node_x->color = RED;

                        while(child->parent != 0)
                                child = child->parent;
                        tree->root = child;

                        lq_rbremedy_double_black(tree, node_x, node_r);
                }
	}
}


struct RbNode*
lq_rbtree_insert(struct RbTree* tree, const_pointer entry, const_pointer data)
{
	assert(tree != 0 && entry != 0);

	struct RbNode* node = tree->root;
	struct RbNode* current_node = 0;

	if(node == 0) {
		tree->root = lq_rbnode_new(BLACK, 0, entry, data);
                tree->nodes += 1;
		return tree->root;
	}

	// general node insertion for binary search tree
	while(current_node == 0) {
		int result = tree->compare(entry, node->entry);

		if(result < 0) {
			if(node->left == 0)
				current_node = node->left = lq_rbnode_new(RED, node, entry, data); 
			else
				node = node->left;
		} else if(result > 0) {
			if(node->right == 0)
				current_node = node->right = lq_rbnode_new(RED, node, entry, data);
			else
				node = node->right;
		} else {
			return 0;
		}
	}

	// repair binary search tree based on red-black tree properties
	lq_rbremedy_double_red(tree, current_node);

        tree->nodes += 1;

	return current_node;
}


static struct RbNode*
lq_rbremove_external_node(struct RbNode* node)
{
	assert(node != 0);

	struct RbNode* child = (node->left == 0) ? node->right : node->left;

        if(RB_PARENT(node) != 0 && RB_PARENT(node)->left == node)
		RB_PARENT(node)->left = child;
	else if(RB_PARENT(node) != 0)
		RB_PARENT(node)->right = child;

	if(child != 0) 
		child->parent = RB_PARENT(node);

	return child;
}


struct RbNode*
lq_rbtree_remove(struct RbTree* tree, const_pointer entry)
{
	assert(tree != 0 && entry != 0);

	struct RbNode* node = 0;
	struct RbNode* current_node = 0;
	struct RbNode* data_node = 0;

	// search for a node with the same entry within the binary tree
	current_node = lq_rbtree_lookup(tree, entry); 

	if(current_node == 0)
		return current_node;
	else if(current_node->left != 0 && current_node->right != 0) {
		// Find most-left internal node with an external child in the right subtree
		node = current_node->right;

		while(node->left != 0)
			node = node->left;

		// swap its entry information
		const_pointer tmp_entry = node->entry;
                const_pointer tmp_data = node->data;
		node->entry = current_node->entry;
                node->data = current_node->data;
		current_node->entry = tmp_entry;
                current_node->data = tmp_data;

		data_node = node;
		node = lq_rbremove_external_node(node);
	} else {
		// Remove internal node and replace it by its single subtree (Easy case)
		data_node = current_node;
		node = lq_rbremove_external_node(current_node);
	}

        // Node removed, now correct tree for red_black_tree properties
        if(RB_PARENT(data_node) == 0) {
                if(node != 0) 
                        node->color = BLACK;
                tree->root = node;
        } else if(node != 0 && (node->color == RED || RB_PARENT(data_node)->color == RED)) {
                node->color = BLACK;
        } else {
                lq_rbremedy_double_black(tree, RB_PARENT(data_node), node);
        }

	tree->nodes -= 1;

	return data_node;
}


struct RbNode*          
lq_rbnode_first(struct RbTree* tree)
{
	assert(tree != 0);

	if(tree->root == 0) 
		return 0;

	struct RbNode* node = tree->root;

	while(node->left != 0) {
		node = node->left;
	}

	return node;
}


struct RbNode*          
lq_rbnode_last(struct RbTree* tree)
{
	assert(tree != 0);

	if(tree->root == 0) 
		return 0;

	struct RbNode* node = tree->root;

	while(node->right != 0) {
		node = node->right;
	}

	return node;
}



struct RbNode*          
lq_rbnode_next(struct RbNode* node)
{
	assert(node != 0);

	struct RbNode* tmp = 0;

	if(node->right != 0) {
		tmp = node->right;

		while(tmp->left != 0)
			tmp = tmp->left;

		return tmp;
	}

	tmp = node->parent;

	while(tmp != 0 && node == tmp->right) {
		node = tmp;
		tmp = tmp->parent; 
	}

	return tmp;
}


struct RbNode*          
lq_rbnode_prev(struct RbNode* node)
{
	assert(node != 0);

	struct RbNode* tmp = 0;

	if(node->left != 0) {
		tmp = node->left;

		while(tmp->right != 0)
			tmp = tmp->right;

		return tmp;
	}

	tmp = node->parent;

	while(tmp != 0 && node == tmp->left) {
		node = tmp;
		tmp = tmp->parent; 
	}

	return tmp;
}


